package com.easy.framework.web.safe.advice;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.easy.framework.common.config.SecretProperties;
import com.easy.framework.common.model.R;
import com.easy.framework.common.utils.Sm2Utils;
import com.easy.framework.web.safe.annotation.ResponseEncrypt;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

@Slf4j
@ControllerAdvice
public class EncryptResponseBodyAdvice implements ResponseBodyAdvice<Object> {

    private boolean encrypt;

    @Autowired
    private SecretProperties secretProperties;

    private static ThreadLocal<Boolean> encryptLocal = new ThreadLocal<>();

    @Override
    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        encrypt = false;
        if (returnType.getMethod().isAnnotationPresent(ResponseEncrypt.class) && !secretProperties.isDebug()) {
            encrypt = true;
        }
        return encrypt;
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType,
                                  Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        // EncryptResponseBodyAdvice.setEncryptStatus(false);
        // Dynamic Settings Not Encrypted
        Boolean status = encryptLocal.get();
        if (null != status && !status) {
            encryptLocal.remove();
            return body;
        }
        if (encrypt) {
            String publicKey = secretProperties.getPublicKey();
            try {
                long l = System.currentTimeMillis();
                if (body instanceof R) {
                    R r = (R) body;
                    if (null != r.getData()) {
                        String data = JSON.toJSONString(r.getData(), SerializerFeature.WriteMapNullValue, SerializerFeature.WriteDateUseDateFormat);
                        r.setData(Sm2Utils.encrypt(publicKey, data));
                    }
                    System.out.println("-----------R 耗时--------------" + (System.currentTimeMillis() - l));
                    return r;
                } else {
                    String content = JSON.toJSONString(body, SerializerFeature.WriteMapNullValue, SerializerFeature.WriteDateUseDateFormat);
                    String result = Sm2Utils.encrypt(publicKey, content);
                    System.out.println("-----------content 耗时--------------" + (System.currentTimeMillis() - l));
                    return result;
                }
            } catch (Exception e) {
                throw new IllegalArgumentException("加密数据异常!");
            }
        }
        return body;
    }
}
